# 这里就是直接调用IK，传参是mocap点的位置和当前关节空间位置
# 然后把结果在mujoco里面显示出来
import socket
import threading
from collections import deque
import struct
import mujoco
import mujoco.viewer as viewer
import numpy as np
from os.path import dirname, join
import time


# 导入自定义函数和CANFD SDK动态链接库>>>
from os.path import dirname, join
import sys
import platform
arch = platform.machine()
sys.path.append(join(dirname(__file__), 'src/lysdemo_'+arch))
sys.path.append(join(dirname(__file__), "src"))
sys.path.append(dirname(__file__))
from src import *



# 共享数据缓冲区
class SharedBuffer:
    def __init__(self):
        self.buffers = [deque(maxlen=1), deque(maxlen=1)]  # 双缓冲
        self.current = 0
        self.lock = threading.Lock()
    
    def write(self, data):
        with self.lock:
            self.buffers[1 - self.current].append(data)
    
    def read(self):
        with self.lock:
            self.current = 1 - self.current
        return self.buffers[self.current][-1] if self.buffers[self.current] else None

# 初始化缓冲区
listA_buffer = SharedBuffer()
listB_buffer = SharedBuffer()

# UDP配置
UDP_IP_B = "192.168.10.13"  # 需替换实际IP
UDP_PORT_SEND = 10000  # A的发送端口
UDP_PORT_RECV = 10001  # B的发送端口

# 发送线程（A->B）
def sender():
    sock_send = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    while True:
        data = listA_buffer.read()
        if data:
            # 二进制打包数据（根据实际数据类型调整格式）
            packed = struct.pack('!I' + 'f'*len(data), len(data), *data)
            sock_send.sendto(packed, (UDP_IP_B, UDP_PORT_SEND))
        time.sleep(0.002)  # 500Hz

# 接收线程（B->A）
def receiver():
    sock_recv = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock_recv.bind(("0.0.0.0", UDP_PORT_RECV))
    while True:
        data, _ = sock_recv.recvfrom(1024)
        # 二进制解包（与B端格式对应）
        length = struct.unpack('!I', data[:4])[0]
        listB = struct.unpack('!' + 'f'*length, data[4:])
        listB_buffer.write(listB)



pos_sensor_names = [
        "Motor_Shoulder_pos", 
        "Motor_BigArm1_pos",
        "Motor_BigArm2_pos",
        "Motor_SmallArm1_pos",
        "Motor_SmallArm2_pos",
        "Motor_Hand_pos"
        ]

vel_sensor_names = [
        "Motor_Shoulder_vel", 
        "Motor_BigArm1_vel",
        "Motor_BigArm2_vel",
        "Motor_SmallArm1_vel",
        "Motor_SmallArm2_vel",
        "Motor_Hand_vel"
        ]

torque_sensor_names = [
    "Motor_Shoulder_Torque",
    "Motor_BigArm1_Torque", 
    "Motor_BigArm2_Torque",
    "Motor_SmallArm1_Torque",
    "Motor_SmallArm2_Torque"
]


if __name__ == "__main__":
    # 启动服务
    threading.Thread(target=sender, daemon=True).start()
    threading.Thread(target=receiver, daemon=True).start()


    # mujoco加载模型
    model = mujoco.MjModel.from_xml_path(join(dirname(__file__),'scene_P.xml')) 
    data = mujoco.MjData(model)
    mocap_id = model.body("target").mocapid[0]

    with viewer.launch_passive(model,data) as viewer:
        # Initialize the camera view to that of the free camera.
        mujoco.mjv_defaultFreeCamera(model, viewer.cam)
        viewer.opt.frame = mujoco.mjtFrame.mjFRAME_SITE

        while viewer.is_running():
            q_real = np.array([data.sensor(pos_sensor_name).data[0] for pos_sensor_name in pos_sensor_names])

            mocap_pos = data.mocap_pos[mocap_id]

            mocap_quat = data.mocap_quat[mocap_id]

            listA_buffer.write(np.concatenate((mocap_pos, mocap_quat)).tolist())


            # 读取listB（非阻塞）
            current_listB = listB_buffer.read()
            if current_listB:
                # 处理listB逻辑
                print(f"current_listB: {current_listB}")
                data.qpos = np.array(current_listB)# 展示ik位置
            
            mujoco.mj_step(model, data)
            mujoco.mj_forward(model, data)
            

            viewer.sync()
            time.sleep(model.opt.timestep)